package cn.itcast.flink.broadcast;

import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.util.Collector;

import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * Author itcast
 * Date 2022/1/16 15:16
 * Desc 实现在socket端输入广告id，实时查询到广告位的信息（广告信息，广告有效期，广告位置等）
 */
public class BroadcastDemo {
    public static void main(String[] args) throws Exception {
        //实现步骤；
        //1）初始化flink流式处理的运行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        //2）设置checkpoint周期性运行，每5秒钟运行一次
        env.enableCheckpointing(5000);
        //3）构建数据源（两个数据源）
        DataStreamSource<String> socket = env.socketTextStream("node1", 9999);
        //构建一个广告数据流
        DataStreamSource<Map<Integer, Tuple2<String, String>>> advertiseStream = env
                .addSource(new MySourceForBroadcastFunction());
        //定义一个广播状态的描述器
        MapStateDescriptor<Integer,Tuple2<String,String>> mapStateDesc = new MapStateDescriptor(
                "advertiseState",
                Types.INT,
                Types.TUPLE(Types.STRING, Types.STRING)
        );
        //4）将广告流作为配置流转换成广播流
        BroadcastStream<Map<Integer, Tuple2<String, String>>> broadcastStream = advertiseStream.broadcast(
                mapStateDesc
        );
        //5）将两个流进行关联操作（connect）
        SingleOutputStreamOperator<Tuple2<String, String>> processStream = socket.connect(broadcastStream)
                //6）对关联的数据进行拉宽操作
                .process(new BroadcastProcessFunction<String, Map<Integer, Tuple2<String, String>>, Tuple2<String, String>>() {
                    @Override
                    public void processElement(String value, ReadOnlyContext ctx, Collector<Tuple2<String, String>> out) throws Exception {
                        //拿到当前内存中的广播状态
                        ReadOnlyBroadcastState<Integer, Tuple2<String, String>> broadcastState = ctx.getBroadcastState(mapStateDesc);
                        //根据广告id将广告信息取出来
                        Tuple2<String, String> advertise = broadcastState.get(Integer.parseInt(value));
                        //将广告id对应的广播信息发送出去
                        out.collect(advertise);
                    }

                    @Override
                    public void processBroadcastElement(Map<Integer, Tuple2<String, String>> value, Context ctx, Collector<Tuple2<String, String>> out) throws Exception {
                        //获取广告状态数据
                        BroadcastState<Integer, Tuple2<String, String>> broadcastState = ctx.getBroadcastState(mapStateDesc);
                        //将最新的广告配置数据更新到状态中
                        broadcastState.clear();
                        broadcastState.putAll(value);
                    }
                });
        //7）打印输出
        processStream.print();
        //8）运行作业
        env.execute();
    }

    /**
     * 广告位流
     */
    public static class MySourceForBroadcastFunction implements SourceFunction<Map<Integer, Tuple2<String, String>>> {
        private final Random random = new Random();
        private final List<Tuple2<String, String>> ads = Arrays.asList(
                Tuple2.of("baidu", "搜索引擎"),
                Tuple2.of("google", "科技大牛"),
                Tuple2.of("aws", "全球领先的云平台"),
                Tuple2.of("aliyun", "全球领先的云平台"),
                Tuple2.of("腾讯", "氪金使我变强"),
                Tuple2.of("阿里巴巴", "电商龙头"),
                Tuple2.of("字节跳动", "靠算法出名"),
                Tuple2.of("美团", "黄色小公司"),
                Tuple2.of("饿了么", "蓝色小公司"),
                Tuple2.of("瑞幸咖啡", "就是好喝")
        );
        private boolean isRun = true;

        @Override
        public void run(SourceContext<Map<Integer, Tuple2<String, String>>> ctx) throws Exception {
            while (isRun) {
                Map<Integer, Tuple2<String, String>> map = new HashMap<>();
                int keyCounter = 0;
                for (int i = 0; i < ads.size(); i++) {
                    keyCounter++;
                    map.put(keyCounter, ads.get(random.nextInt(ads.size())));
                }
                ctx.collect(map);
                TimeUnit.SECONDS.sleep(5L);
            }
        }

        @Override
        public void cancel() {
            this.isRun = false;
        }
    }
}
